package com.hivclassifier.gui;

import java.awt.BorderLayout;
import java.awt.Dimension;
import java.awt.GridBagConstraints;
import java.awt.GridBagLayout;
import java.awt.Insets;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.io.File;

import javax.swing.JButton;
import javax.swing.JDialog;
import javax.swing.JFileChooser;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JOptionPane;
import javax.swing.JPanel;
import javax.swing.JProgressBar;
import javax.swing.JScrollPane;
import javax.swing.JTextArea;
import javax.swing.ScrollPaneConstants;

import weka.classifiers.Classifier;
import weka.classifiers.Evaluation;

import com.hivclassifier.HivClassifierController;
import com.hivclassifier.utils.Logger;

public class MainScreen extends JFrame {

	private static final long serialVersionUID = -4263002646581609576L;

	private JTextArea log;

	private JButton btnLoadClassifier = new JButton("Load classifier");
	private JButton btnSaveClassifier = new JButton("Save classifier");
	private JButton btnAppendToTrainingSet = new JButton("Apply data to classifier");
	private JButton btnPredictCoreceptorAffinity = new JButton("Predict coreceptor affinity");
	private JButton btnNewClassifier = new JButton("New classifier");
	
	public MainScreen() {

		setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);

		setTitle("B4-HIV-1 predictor");

		JPanel rootPanel = new JPanel();
		rootPanel.setPreferredSize(new Dimension(800, 600));

		GridBagLayout gridBagLayout = new GridBagLayout();
		gridBagLayout.columnWidths = new int[] { 0 };
		gridBagLayout.columnWeights = new double[] { 1 };
		gridBagLayout.rowHeights = new int[] { 0, 0, };
		gridBagLayout.rowWeights = new double[] { 0.2, 0.8, };
		rootPanel.setLayout(gridBagLayout);

		gridBagLayout = new GridBagLayout();
		gridBagLayout.columnWidths = new int[] { 0, 0, 0, 0 };
		gridBagLayout.columnWeights = new double[] { 0.20, 0.20, 0.20, 0.20, 0.20};
		gridBagLayout.rowHeights = new int[] { 0, };
		gridBagLayout.rowWeights = new double[] { 1, };

		JPanel panelMainButtons = new JPanel(gridBagLayout);
		GridBagConstraints gbc = new GridBagConstraints();
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.fill = GridBagConstraints.VERTICAL;
		gbc.gridx = 0;
		gbc.gridy = 0;
		rootPanel.add(panelMainButtons, gbc);

		btnNewClassifier.addActionListener(new ActionListener() {

			@Override
			public void actionPerformed(ActionEvent e) {
				
				saveClassifierDialog();

				HivClassifierController.setBayesClassifier(null);

				btnSaveClassifier.setEnabled(false);
				btnPredictCoreceptorAffinity.setEnabled(false);
			}
		});
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.insets = new Insets(10, 5, 5, 0);
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 0;
		gbc.gridy = 0;
		panelMainButtons.add(btnNewClassifier, gbc);
		
		btnLoadClassifier.addActionListener(new ActionListener() {

			@Override
			public void actionPerformed(ActionEvent e) {
				final JFileChooser fc = new JFileChooser();

				int returnVal = fc.showOpenDialog(MainScreen.this);

				if (returnVal == JFileChooser.APPROVE_OPTION) {
					try {
						Classifier loadedBayesClassifier = (Classifier) weka.core.SerializationHelper
								.read(fc.getSelectedFile().getAbsolutePath());
						HivClassifierController.setBayesClassifier(loadedBayesClassifier);

						Logger.getInstance().doLog(
								"\n Load classifier: " + fc.getSelectedFile().getAbsolutePath());

						if (!btnSaveClassifier.isEnabled()) {
							btnSaveClassifier.setEnabled(true);
						}
						if (!btnPredictCoreceptorAffinity.isEnabled()) {
							btnPredictCoreceptorAffinity.setEnabled(true);
						}
					} catch (Exception e1) {
						JOptionPane.showMessageDialog(MainScreen.this,
								"Unable to load classifier!", "Classifier loading error",
								JOptionPane.ERROR_MESSAGE);
					}
				}
			}
		});
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.insets = new Insets(10, 5, 5, 0);
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 1;
		gbc.gridy = 0;
		panelMainButtons.add(btnLoadClassifier, gbc);

		btnSaveClassifier.setEnabled(false);
		btnSaveClassifier.addActionListener(new ActionListener() {

			@Override
			public void actionPerformed(ActionEvent e) {
				final JFileChooser fc = new JFileChooser();

				int returnVal = fc.showOpenDialog(MainScreen.this);

				if (returnVal == JFileChooser.APPROVE_OPTION) {
					File file = fc.getSelectedFile();
					Logger.getInstance().doLog(
							"Saving classifier to: " + file.getAbsolutePath() + "\n");
					try {
						weka.core.SerializationHelper.write(file.getAbsolutePath(),
								HivClassifierController.getBayesClassifier());
					} catch (Exception ex) {
						ex.printStackTrace();
					}
				}
			}
		});
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.insets = new Insets(10, 5, 5, 0);
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 2;
		gbc.gridy = 0;
		panelMainButtons.add(btnSaveClassifier, gbc);

		btnAppendToTrainingSet.addActionListener(new ActionListener() {

			@Override
			public void actionPerformed(ActionEvent e) {

				final JDialog dlg = new JDialog(MainScreen.this, "Please wait", true);
				dlg.setDefaultCloseOperation(JDialog.DO_NOTHING_ON_CLOSE);
				dlg.setSize(300, 75);
				dlg.setLocationRelativeTo(MainScreen.this);

				final JFileChooser fc = new JFileChooser();

				int returnVal = fc.showOpenDialog(MainScreen.this);

				if (returnVal == JFileChooser.APPROVE_OPTION) {

					new Thread(new Runnable() {
						public void run() {
							
							JProgressBar dpb = new JProgressBar();
							
							dlg.add(BorderLayout.CENTER, dpb);
							dlg.add(BorderLayout.NORTH, new JLabel("Training..."));
							
							new Thread(new Runnable() {
								public void run() {
									dlg.setVisible(true);
								}
							}).start();

							dpb.setIndeterminate(true);
							
							Evaluation evaluation = HivClassifierController
									.appendToTrainingSetAndEvaluate(fc.getSelectedFile().getAbsolutePath());
							
							dlg.setVisible(false);
							dlg.dispose();
							
							if (evaluation != null) {
								new StatisticsResultScreen(evaluation);
								Logger.getInstance().doLog(
										"\n Appending to training set file: "
												+ fc.getSelectedFile().getAbsolutePath());

								if (!btnSaveClassifier.isEnabled()) {
									btnSaveClassifier.setEnabled(true);
								}
								if (!btnPredictCoreceptorAffinity.isEnabled()) {
									btnPredictCoreceptorAffinity.setEnabled(true);
								}
							} else {
								JOptionPane.showMessageDialog(MainScreen.this,
										"Unable to train classifier with this file!", "Classifier training error",
										JOptionPane.ERROR_MESSAGE);
							}
						}
					}).start();
				}
			}
		});
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.insets = new Insets(10, 5, 5, 0);
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 3;
		gbc.gridy = 0;
		panelMainButtons.add(btnAppendToTrainingSet, gbc);

		btnPredictCoreceptorAffinity.setEnabled(false);
		btnPredictCoreceptorAffinity.addActionListener(new ActionListener() {

			@Override
			public void actionPerformed(ActionEvent e) {

				final JDialog dlg = new JDialog(MainScreen.this, "Please wait", true);
				dlg.setDefaultCloseOperation(JDialog.DO_NOTHING_ON_CLOSE);
				dlg.setSize(300, 75);
				dlg.setLocationRelativeTo(MainScreen.this);
				
				final JFileChooser fc = new JFileChooser();

				int returnVal = fc.showOpenDialog(MainScreen.this);

				if (returnVal == JFileChooser.APPROVE_OPTION) {
					// TODO make a spinner dialog to keep the user entertained and
					// do in background
					
					new Thread(new Runnable() {
						public void run() {
							
							JProgressBar dpb = new JProgressBar();
							
							dlg.add(BorderLayout.CENTER, dpb);
							dlg.add(BorderLayout.NORTH, new JLabel("Classifying..."));
							
							new Thread(new Runnable() {
								public void run() {
									dlg.setVisible(true);
								}
							}).start();

							dpb.setIndeterminate(true);
							
							String[][] result = HivClassifierController.performClassification(fc
									.getSelectedFile().getAbsolutePath());
							
							dlg.setVisible(false);
							dlg.dispose();
							
							if (result != null) {
								new PredictionsResultScreen(result);
								
								Logger.getInstance().doLog(
										"\n Predicting coreceptor affinity for file: "
												+ fc.getSelectedFile().getAbsolutePath());
							} else {
								JOptionPane.showMessageDialog(MainScreen.this,
										"Unable to perform classification on this file!", "Classifier prediction error",
										JOptionPane.ERROR_MESSAGE);
							}
						}
					}).start();
					
				}
			}
		});
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.insets = new Insets(10, 5, 5, 0);
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 4;
		gbc.gridy = 0;
		panelMainButtons.add(btnPredictCoreceptorAffinity, gbc);

		JPanel panelStatus = new JPanel();
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 0;
		gbc.gridy = 1;
		rootPanel.add(panelStatus, gbc);

		gridBagLayout = new GridBagLayout();
		gridBagLayout.columnWidths = new int[] { 0 };
		gridBagLayout.columnWeights = new double[] { 1 };
		gridBagLayout.rowHeights = new int[] { 0, 0, };
		gridBagLayout.rowWeights = new double[] { 0.1, 0.9, };
		panelStatus.setLayout(gridBagLayout);

		gbc.anchor = GridBagConstraints.WEST;
		gbc.fill = GridBagConstraints.VERTICAL;
		gbc.insets = new Insets(0, 10, 5, 10);
		gbc.gridx = 0;
		gbc.gridy = 0;
		panelStatus.add(new JLabel("Log:"), gbc);

		log = new JTextArea(20, 55);
		log.setEditable(false);
		JScrollPane scroll = new JScrollPane(log);
		scroll.setVerticalScrollBarPolicy(ScrollPaneConstants.VERTICAL_SCROLLBAR_ALWAYS);
		gbc.anchor = GridBagConstraints.CENTER;
		gbc.fill = GridBagConstraints.BOTH;
		gbc.gridx = 0;
		gbc.gridy = 1;

		panelStatus.add(scroll, gbc);

		// Display the window.
		this.getContentPane().add(rootPanel);
		this.setVisible(true);
		this.pack();

		addWindowListener(new java.awt.event.WindowAdapter() {
			@Override
			public void windowClosing(java.awt.event.WindowEvent windowEvent) {

				saveClassifierDialog();

				System.exit(0);
			}
		});
	}
	
	public void saveClassifierDialog() {
		
		if (HivClassifierController.getBayesClassifier() != null) {
			boolean exit = false;
			
			while (!exit) {
				
				if (JOptionPane.showConfirmDialog(MainScreen.this,
						"Would you like to save the current classifier model?",
						"Closing HIV-1 correceptor affinity", JOptionPane.YES_NO_OPTION,
						JOptionPane.QUESTION_MESSAGE) == JOptionPane.YES_OPTION) {
					
					final JFileChooser fc = new JFileChooser();
					
					int returnVal = fc.showSaveDialog(MainScreen.this);
					
					if (returnVal == JFileChooser.APPROVE_OPTION) {
						File file = fc.getSelectedFile();
						try {
							weka.core.SerializationHelper.write(file.getAbsolutePath(),
									HivClassifierController.getBayesClassifier());
						} catch (Exception e) {
							e.printStackTrace();
						}
						
						exit = true;
					}
				} else {
					exit = true;
				}
			}
		}
	}

	public JTextArea getLogTextField() {
		return log;
	}
}
